#include <iostream>
#include <fstream>
#include <shared_mutex>
#include <vector>
#include <bitset>
using namespace std;

struct HuffmanTreeNode
{
    HuffmanTreeNode(char ch, int weight=0, int lchild=-1, int rchild=-1, int parent=-1)
     : ch(ch), weight(weight), lchild(lchild), rchild(rchild), parent(parent) {}
    char ch;
    int weight;
    int lchild;
    int rchild;
    int parent;
};

struct CharCounter
{
    CharCounter(char ch, int cnt):ch(ch), counter(cnt){}
    char ch;
    int counter=0;
};
// 由于无法直接实现位级别的动态存储，故而采用了索引的方式
struct HuffmanCode
{
    HuffmanCode(char ch='\0') : ch(ch){}
    char ch;
    unsigned char start_index=0;
    std::bitset<8> code;
};

class HuffmanTree
{
public:
    bool readText(std::string filepath)
    {
        read_file.open(filepath, ios::in);
        if(!read_file.is_open())
        {
            return false;
        }
        size_t length=0;
        while(!read_file.eof())
        {
            read_file >> buff;
            buff.append("\n");
            generateCounts(buff);
            length += buff.size();
            buff.clear();
        }
        buff = filepath;
        cout << "original file size: " << length << "Bytes" << endl;
        read_file.close();
        n = counts.size();
        for(size_t i = 0; i < 2*n-1; ++i)
        {
            if(i < n)
                HT.emplace_back(counts[i].ch, counts[i].counter);
            else
                HT.emplace_back('\0', INT_MAX);
        }
        createHT();
        return true;
    }
    
    std::vector<HuffmanCode> encode()
    {
        HuffmanCode temp_code;
        int parent_pos, child_pos, cnt;
        for(size_t i = 0; i < n; ++i)
        {
            cnt = 7;
            temp_code.ch = HT[i].ch;
            child_pos = i;
            while((parent_pos=HT[child_pos].parent)>=0)
            {
                temp_code.code[cnt--] = (HT[parent_pos].lchild==child_pos)?0:1;
                child_pos = parent_pos;
            }
            temp_code.start_index = cnt+1;
            HC.emplace_back(temp_code);
        }
        return HC;
    }
    // 实现了二进制形式存储
    bool zipFile(std::string filepath)
    {
        std::shared_ptr<std::bitset<256>> bits(new std::bitset<256>);
        unsigned short cnt=0;
        size_t length=0;
        write_file.open(filepath, ios::binary | ios::trunc);
        if(write_file.is_open())
        {
            read_file.open(buff, ios::in);
            while(!read_file.eof())
            {
                read_file >> buff;
                buff.append("\n");
                for(size_t i = 0; i < buff.size(); ++i)
                {
                    for(size_t j = 0; j < HC.size(); ++j)
                    {
                        if(HC[j].ch == buff.at(i))
                        {
                            for(unsigned char k = HC[j].start_index; k < 8; ++k)
                            {
                                bits->set(cnt++, HC[j].code[k]);
                                if(cnt == 256)
                                {
                                    length += cnt;
                                    write_file.write((char*)bits.get(), 32);
                                    cnt = 0;
                                }
                            }
                        }
                    }
                }
                buff.clear();
            }
            length += cnt;
            // 多预留一个字节，否则将缺失一块内容
            write_file.write((char*)bits.get(), cnt%8==0?cnt/8:cnt/8+1);
            cout << "zipped file size: " << length/8 << "Bytes" << endl;
            read_file.close();
            write_file.close();
            return true;
        }
        else
            return false;
    }

    // 二进制解码，由于在获取二进制块并转移的过程中的边界条件太多，暂时没有做完
    // 有另一种思路：先把所有数据读取完，压入vector后再进行解码操作
    std::string decode(std::string filepath)
    {
        read_file.open(filepath, ios::binary);
        if(!read_file.is_open())
        {
            cerr << "Cannot open file: " << filepath << endl;
            return "";
        }
        else
        {
            bool correct;
            unsigned char cnt1=0;
            int byte_num, bound=0;
            size_t cnt2=0;
            std::string result;
            std::shared_ptr<std::bitset<256>> bits(new std::bitset<256>), temp_bits(new std::bitset<256>);
            while(!read_file.eof())
            {
                read_file.read((char*)bits.get(), 32);
                // 获取当前的字节数
                byte_num = read_file.gcount();
                if(byte_num == 32)
                    bitsets.push_back(bits);
                else
                {
                    *temp_bits = 0;
                    bitsets.push_back(temp_bits);
                    // *(*bitsets.end()) = 0;
                    for(short i = 0; i < byte_num*8; ++i)
                        (*(*(bitsets.end()-1)))[i] = (*bits)[i]; 
                }
            }
            // cout << bitsets.size() << endl;
            // cout << (bitsets.size()-1)*256 << endl;
            // cout << (bitsets.size()-1)*256+byte_num*8 << endl;
            size_t max_size = (bitsets.size()-1)*256+byte_num*8;
            while(cnt2 < max_size)
            {
                for(int i = 0; i < HC.size(); ++i)
                {
                    bound = 0;
                    cnt1 = HC[i].start_index;
                    correct = false;
                    while((*bits)[cnt2++] == HC[i].code[cnt1++])
                    {
                        if(cnt1 == 8)
                        {
                            correct = true;
                            result.push_back(HC[i].ch);
                            break;
                        }
                        if(cnt2 >= max_size)
                            break;
                    }
                    if(cnt2 >= max_size)
                        break;
                    if(!correct)
                    {
                        cnt2 -= cnt1-HC[i].start_index;
                    }
                }
                if(cnt2 == max_size-(cnt1-HC[HC.size()-1].start_index))
                {
                    cerr << "Err: The file is not Huffman format binary code!" << endl;
                    return "";
                }
            }
            read_file.close();
            return result;
        }
    }
private:
    void inline createHT()
    {
        int p1, p2;
        for(size_t i=n; i < 2*n-1; ++i)
        {
            findMinNode(p1, p2);
            HT[p1].parent=i;
            HT[p2].parent=i;
            HT[i].lchild = p1;
            HT[i].rchild = p2;
            HT[i].weight = HT[p1].weight+HT[p2].weight;
        }
    }
    void inline findMinNode(int& least_1, int& least_2)
    {
        least_1 = 0;
        least_2 = 1;
        int min_value=INT_MAX, min_value2=INT_MAX;
        for(size_t i=0; i < 2*n-1; ++i)
        {
            if(HT[i].parent == -1)
            {
                if(HT[i].weight < min_value)
                {
                    min_value = HT[i].weight;
                    least_1 = i;
                }
                else if(HT[i].weight < min_value2)
                {
                    min_value2 = HT[i].weight;
                    least_2 = i;
                }
            }
        }
    }
    void inline generateCounts(std::string& buff)
    {
        bool is_new = true;
        for(size_t i = 0; i < buff.size(); ++i)
        {
            for(auto iter = counts.begin(); iter != counts.end(); iter++)
            {
                if(buff.at(i) == (*iter).ch)
                {
                    (*iter).counter++;
                    is_new = false;
                    break;
                }
            }
            if(is_new == true)
            {
                counts.emplace_back(buff.at(i), 1);
            }
            else
                is_new = true;
        }
    }
    std::vector<CharCounter> counts;
    std::vector<std::shared_ptr<std::bitset<256>>> bitsets;
    std::vector<HuffmanTreeNode> HT;
    std::vector<HuffmanCode> HC;
    std::string buff;
    ifstream read_file;
    ofstream write_file;
    size_t n;
};

int main(int argc, const char** argv) 
{
    if(argc != 3)
    {
        cerr << "usage: huffman " << "<input_file>" << "<output_file>" << endl;
        return 1;
    }
    HuffmanTree T;
    T.readText(argv[1]);
    auto result = T.encode();
    T.zipFile(argv[2]);
    std::string decode_result = T.decode(argv[2]);
    for(auto i : result)
    {
        if(i.ch == '\n')
            cout << "\\n" << ": ";
        else
            cout << i.ch << ": ";
        for(unsigned char j = i.start_index; j < 8; ++j)
            cout << i.code[j];
        cout << endl;
    }
    cout << "The decode result is:\n" << decode_result << endl;
    return 0;
}